import os
import random
import glob

# 设置随机种子以确保可重复性
random.seed(42)

# 定义数据集路径
dataset_path = os.path.expanduser("~/yolov12_workspace/armor_dataset")
images_dir = os.path.join(dataset_path, "images")

# 获取所有图像文件
image_files = glob.glob(os.path.join(images_dir, "*.jpg"))
image_files.sort()  # 确保文件排序

# 随机打乱文件顺序用于分割
file_indices = list(range(len(image_files)))
random.shuffle(file_indices)

# 分割为训练集和验证集（默认80%训练，20%验证）
train_ratio = 0.8
split_index = int(len(file_indices) * train_ratio)
train_indices = file_indices[:split_index]
val_indices = file_indices[split_index:]

# 创建相对路径列表
train_paths = []
val_paths = []

for idx in train_indices:
    image_path = os.path.relpath(image_files[idx], dataset_path)
    train_paths.append(image_path)

for idx in val_indices:
    image_path = os.path.relpath(image_files[idx], dataset_path)
    val_paths.append(image_path)

# 保存训练集文件列表
train_file_path = os.path.join(dataset_path, "train.txt")
with open(train_file_path, "w") as f:
    f.write("\n".join(train_paths))

# 保存验证集文件列表
val_file_path = os.path.join(dataset_path, "val.txt")
with open(val_file_path, "w") as f:
    f.write("\n".join(val_paths))

print(f"总图像数量: {len(image_files)}")
print(f"训练集图像数量: {len(train_paths)}")
print(f"验证集图像数量: {len(val_paths)}")
print(f"训练集文件已保存到: {train_file_path}")
print(f"验证集文件已保存到: {val_file_path}")

# 修改rune.yaml文件中的路径
rune_yaml_path = "/home/ccy/Downloads/2/1/rune.yaml"
with open(rune_yaml_path, "r") as f:
    lines = f.readlines()

# 更新路径
for i, line in enumerate(lines):
    if line.startswith("path:"):
        lines[i] = f"path: {dataset_path}  # 数据集路径\n"

# 保存修改后的yaml文件
with open(rune_yaml_path, "w") as f:
    f.writelines(lines)

print(f"已更新数据集路径在: {rune_yaml_path}")